/*
 * Copyright (C) 2024 KylinSoft Co., Ltd.
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later
 * version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program. If not, see <https://www.gnu.org/licenses/>.
 */

#include "rpc/vectorDB/vectorDatabase.h"

#include "log.h"

using namespace scene;

VectorDatabase& VectorDatabase::getInstance() {
    static VectorDatabase vectorDB;
    return vectorDB;
}

VectorDatabase::VectorDatabase() : isConnected(false) { client = VectorDB::Database::Create(); }

bool VectorDatabase::connect() {
    mtx.lock();
    if (isConnected) {
        mtx.unlock();
        return true;
    }

    auto status = client->Connect();
    if (status.IsOk()) {
        isConnected = true;
        mtx.unlock();
        return true;
    }
    mtx.unlock();
    return checkStatus("Failed to connect vector db", status);
}

bool VectorDatabase::getConnection() {
    mtx.lock_shared();
    if (isConnected) {
        mtx.unlock_shared();
        return true;
    }
    mtx.unlock_shared();

    return connect();
}

bool VectorDatabase::checkStatus(std::string&& prefix, const VectorDB::Status& status) {
    if (!status.IsOk()) {
        PRINT_ERROR("%s, code: %d, message: %s\n", prefix.c_str(), static_cast<int>(status.Code()),
                    status.Message().c_str());
        return false;
    }
    return true;
}

bool VectorDatabase::hasCollection(const std::string& name) {
    GET_CONNECTION()

    bool exist;
    auto status = client->HasCollection(name, exist);
    // 数据库接口有问题，临时兼容下，如果返回报错也认为不存在，Failed to has TEST collection: 1002: collection not found
    if (!status.IsOk()) {
        return false;
    }
    return exist;
}

bool VectorDatabase::createCollection(const std::string& name, int dimension, bool autoId, bool enableDynamicField) {
    GET_CONNECTION()

    if (hasCollection(name)) {
        return true;
    }

    auto status = client->CreateCollection(name, dimension, autoId, enableDynamicField);
    return checkStatus("Failed to create collection " + name, status);
}

bool VectorDatabase::createCollection(const VectorDB::CollectionSchema& schema, const VectorDB::IndexDesc& desc) {
    GET_CONNECTION()

    if (hasCollection(schema.Name())) {
        return true;
    }

    auto status = client->CreateCollection(schema, desc);
    return checkStatus("Failed to create collection " + schema.Name(), status);
}

bool VectorDatabase::dropCollection(const std::string& name) {
    GET_CONNECTION()

    if (!hasCollection(name)) {
        return true;
    }

    auto status = client->DropCollection(name);
    return checkStatus("Failed to drop collection " + name, status);
}

bool VectorDatabase::insert(const std::string& name, const std::vector<VectorDB::FieldDataPtr>& fields,
                            VectorDB::DmlResults& results) {
    GET_CONNECTION()

    auto status = client->Insert(name, fields, results);
    return checkStatus("Failed to insert", status);
}

bool VectorDatabase::search(const VectorDB::SearchArguments& arguments, VectorDB::SearchResults& results, int timeout) {
    GET_CONNECTION()

    auto status = client->Search(arguments, results, timeout);
    return checkStatus("Failed to search", status);
}

bool VectorDatabase::query(const VectorDB::QueryArguments& arguments, VectorDB::QueryResults& results, int timeout) {
    GET_CONNECTION()

    auto status = client->Query(arguments, results, timeout);
    return checkStatus("Failed to query", status);
}

bool VectorDatabase::upsert(const std::string& name, const std::vector<VectorDB::FieldDataPtr>& fields,
                            VectorDB::DmlResults& results) {
    GET_CONNECTION()

    auto status = client->Upsert(name, fields, results);
    return checkStatus("Failed to upsert", status);
}

bool VectorDatabase::deleteData(const std::string& name, const std::string& expression, VectorDB::DmlResults& results) {
    GET_CONNECTION()

    auto status = client->Delete(name, expression, results);
    return checkStatus("Failed to delete", status);
}

bool VectorDatabase::disconnect() {
    mtx.lock();
    if (!isConnected) {
        mtx.unlock();
        return true;
    }

    auto status = client->Disconnect();
    if (status.IsOk()) {
        isConnected = false;
        mtx.unlock();
        return true;
    }
    mtx.unlock();
    return checkStatus("Failed to disconnect vector db", status);
}

VectorDatabase::~VectorDatabase() { disconnect(); }
